# 安装依赖 pip install langchain_community
import os

from langchain_openai import AzureChatOpenAI
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

os.environ["AZURE_OPENAI_ENDPOINT"] = "http://menshen.test.xdf.cn"
os.environ["OPENAI_API_KEY"] = "c8575027653b42b1b47747f0b4ab135b"
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"

chat_msg_history = ChatMessageHistory()
chat_msg_history.add_user_message("你好，我叫老大")
chat_msg_history.add_ai_message("你好，我是langchain")
chat_msg_history.add_user_message("我今天下午玩游戏了,心情不错")
chat_msg_history.add_ai_message("好的，请继续")
chat_msg_history.add_user_message("今天下午还学了一些langchain的知识,感觉特别充实")
chat_msg_history.add_ai_message("真替你感到高兴")

llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    model_name="gpt-4o",
    temperature=0
)

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个{role}专家，请用你的专业知识回答问题"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{question}")
])

#  prompt模板链式调用
chain = prompt | llm

# 存储聊天记录
store = {}


def summary_message(chain_input):
    msg_store = chat_msg_history.messages
    if len(msg_store) == 0:
        return False
    summary_prompt = ChatPromptTemplate.from_messages([
        MessagesPlaceholder(variable_name="chat_history"),
        ("user", "请总结当前对话，尽可能包含多个细节,并给出一个总结的句子")
    ])
    summary_chain = summary_prompt | llm
    summary = summary_chain.invoke({"chat_history": msg_store})
    print(f'总结结果：{summary.content}')
    chat_msg_history.clear()
    chat_msg_history.add_message(summary)


do_message = RunnableWithMessageHistory(
    chain,
    lambda session_id: chat_msg_history,
    input_messages_key="question",  # 每次聊天时发送消息的key
    history_messages_key="chat_history",
)

chain_with_trim = (RunnablePassthrough.assign(summary_msg=summary_message) | do_message)

config = {'configurable': {'session_id': 'zs123'}}

# 第1轮
resp = chain_with_trim.invoke(
    {
        "role": "陪聊",
        "question": "我的名字，下午在干嘛，什么心情"
    },
    config
)
print(resp.content)
